package cn.iocoder.yudao.module.infra.vector.impl.dashvector;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.module.infra.vector.IVectorService;
import cn.iocoder.yudao.module.infra.vector.VectorConfig;
import cn.iocoder.yudao.module.resource.api.model.messagecollection.VectorDoc;
import com.aliyun.dashvector.DashVectorClient;
import com.aliyun.dashvector.DashVectorClientConfig;
import com.aliyun.dashvector.DashVectorCollection;
import com.aliyun.dashvector.models.*;
import com.aliyun.dashvector.models.Vector;
import com.aliyun.dashvector.models.requests.DeleteDocRequest;
import com.aliyun.dashvector.models.requests.InsertDocRequest;
import com.aliyun.dashvector.models.requests.QueryDocRequest;
import com.aliyun.dashvector.models.requests.UpdateDocRequest;
import com.aliyun.dashvector.models.responses.Response;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;

import java.util.*;

/**
 * https://help.aliyun.com/document_detail/2837745.html?spm=a2c4g.11186623.help-menu-2510217.d_0_4_8.d9977097pfge7A&scm=20140722.H_2837745._.OR_help-T_cn~zh-V_1
 */
@Slf4j
@Service("vector-dashvector")
public class DashVectorServiceImpl implements IVectorService {
    @Resource
    private VectorConfig vectorConfig;

    private DashVectorClient client;
    private DashVectorCollection collection;
    private VectorConfig.Config dashConfig;

    @PostConstruct
    public void init() {
        if (CollUtil.isEmpty(vectorConfig.getServices())) {
            log.warn("未配置向量存储服务");
            return;
        }

        this.dashConfig = vectorConfig.getService("dashvector");
        if (dashConfig == null) {
            log.warn("未配置dashvector向量存储服务");
            return;
        }

        this.client = newClient();
        this.collection = getCollection(client);
    }

    @Override
    public List<VectorDoc> queryDoc(String partition,
                                    Map<String, List<? extends Number>> vectors,
                                    Map<String, Float> weights,
                                    String filter, Integer topk, Float minSimilarity,
                                    Collection<String> outputFields) {
        log.info("检索doc-开始, partition = {}, filter = {}, minSimilarity = {}, topk = {}, outputFields = {}",
                partition, filter, minSimilarity, topk, outputFields);

        final long startTime = System.currentTimeMillis();
        if (StrUtil.isEmpty(partition)) {
            partition = dashConfig.getPartition();
        }
        if (topk == null) {
            topk = dashConfig.getTopk();
        }
        if (minSimilarity == null) {
            minSimilarity = dashConfig.getMinSimilarity();
        }

        final float radius = 1 - minSimilarity;

        final Map<String, VectorQuery> vectorMap = new HashMap<>();
        vectors.forEach((key, value) -> {
            // 构建Vector
            final Vector queryVector = Vector.builder().value(value).build();
            final VectorQuery vectorQuery = VectorQuery.builder().vector(queryVector).radius(radius).build();
            vectorMap.put(key, vectorQuery);
        });

        // 构建QueryDocRequest
        final QueryDocRequest request = QueryDocRequest.builder()
                .partition(partition)
                .vectors(vectorMap)
                .filter(filter)
                .topk(topk)
                .includeVector(true)
                .ranker(WeightedRanker.builder().weights(weights).build())
                .outputFields(outputFields)
                .build();

        // 进行Doc检索
        final Response<List<Doc>> response = collection.query(request);

        final String requestId = response.getRequestId();
        if (response.isSuccess()) {
            final List<VectorDoc> docs = new ArrayList<>();
            for (Doc doc : response.getOutput()) {
                final float similarity = 1 - doc.getScore();
                log.info("检索doc-过滤, docId = {}, similarity = {}, minSimilarity = {}", doc.getId(), similarity, minSimilarity);
                if (similarity >= minSimilarity) {
                    docs.add(VectorDoc.builder().id(doc.getId()).fields(doc.getFields()).similarity(similarity).build());
                }
            }
            final long duration = System.currentTimeMillis() - startTime;
            log.info("检索doc-成功, duration = {}, requestId = {}, docs = {}", duration, requestId, JsonUtils.toJsonString(docs));
            return docs;
        } else {
            final int code = response.getCode();
            final String message = response.getMessage();
            log.error("检索doc-失败, requestId = {}, code = {}, message = {}",
                    requestId, code, message);
            return null;
        }
    }

    @Override
    public void insertDoc(String docId,
                          Map<String, Object> fields,
                          Map<String, List<? extends Number>> vectorValues) {
        final long startTime = System.currentTimeMillis();
        log.info("插入doc-开始, docId = {}, fields = {}", docId, fields);

        // 构建Vector
        final Map<String, Vector> vectors = buildVectors(vectorValues);

        // 构建Doc
        final Doc doc = Doc.builder().id(docId).vectors(vectors).fields(fields).build();

        // 插入Doc
        final Response<List<DocOpResult>> response = collection.insert(InsertDocRequest.builder().doc(doc).build());

        final String requestId = response.getRequestId();
        if (response.isSuccess()) {
            final long duration = System.currentTimeMillis() - startTime;
            log.info("插入doc-成功, docId = {}, requestId = {}, duration = {}", docId, requestId, duration);
        } else {
            final int code = response.getCode();
            final String message = response.getMessage();
            log.error("插入doc-失败, docId = {}, requestId = {}, code = {}, message = {}",
                    docId, requestId, code, message);
            throw new RuntimeException(StrUtil.format("插入doc-失败, code = {}, message = {}", code, message));
        }
    }

    @Override
    public void updateDoc(String docId,
                          Map<String, Object> fields,
                          Map<String, List<? extends Number>> vectorValues) {
        final long startTime = System.currentTimeMillis();
        log.info("更新doc-开始, docId = {}, fields = {}", docId, fields);

        // 构建Vector
        final Map<String, Vector> vectors = buildVectors(vectorValues);

        // 构建Doc
        final Doc doc = Doc.builder().id(docId).vectors(vectors).fields(fields).build();

        // 插入Doc
        final Response<List<DocOpResult>> response = collection.update(UpdateDocRequest.builder().doc(doc).build());

        final String requestId = response.getRequestId();
        if (response.isSuccess()) {
            final long duration = System.currentTimeMillis() - startTime;
            log.info("更新doc-成功, docId = {}, requestId = {}, duration = {}", docId, requestId, duration);
        } else {
            final int code = response.getCode();
            final String message = response.getMessage();
            log.error("更新doc-失败, docId = {}, requestId = {}, code = {}, message = {}",
                    docId, requestId, code, message);
            throw new RuntimeException(StrUtil.format("更新doc-失败, code = {}, message = {}", code, message));
        }
    }

    @Override
    public void deleteDoc(String docId) {
        final long startTime = System.currentTimeMillis();
        log.info("删除doc-开始, docId = {}", docId);

        final Response<List<DocOpResult>> response = collection.delete(DeleteDocRequest.builder().id(docId).build());
        final String requestId = response.getRequestId();
        if (response.isSuccess()) {
            final long duration = System.currentTimeMillis() - startTime;
            log.info("删除doc-成功, docId = {}, requestId = {}, duration = {}", docId, requestId, duration);
        } else {
            final int code = response.getCode();
            final String message = response.getMessage();
            log.error("删除doc-失败, docId = {}, requestId = {}, code = {}, message = {}",
                    docId, requestId, code, message);
            throw new RuntimeException(StrUtil.format("删除doc-失败, code = {}, message = {}", code, message));
        }
    }

    @NotNull
    private static Map<String, Vector> buildVectors(Map<String, List<? extends Number>> vectorValues) {
        final Map<String, Vector> vectors = new HashMap<>();
        vectorValues.forEach((key, value) -> {
            final Vector vector = Vector.builder().value(value).build();
            vectors.put(key, vector);
        });
        return vectors;
    }

    private DashVectorCollection getCollection(DashVectorClient client) {
        return client.get(dashConfig.getCollection());
    }

    private DashVectorClient newClient() {
        final String apiKey = dashConfig.getApiKey();
        final String apiUrl = dashConfig.getApiUrl();

        // 通过Builder构造DashVectorClientConfig
        final DashVectorClientConfig config = DashVectorClientConfig.builder()
                .apiKey(apiKey)
                .endpoint(apiUrl)
                .timeout(30f)
                .build();
        final DashVectorClient client = new DashVectorClient(config);
        return client;
    }


    public static void main(String[] args) {

    }

}
